namespace NHCodeFirst
{
    using System;
    using System.Collections.Generic;
    using System.Linq.Expressions;
    using NHibernate.Linq;

    public class FetchedDbSet<T, TProperty> : DbSetBase<T,FetchedDbSet<T, TProperty>> where T :class
    {
        private readonly INhFetchRequest<T, TProperty> fetchRequest;

        public FetchedDbSet(DbContext context, INhFetchRequest<T,TProperty> fetchRequest) : base(context)
        {
            this.fetchRequest = fetchRequest;
        }

        public FetchedDbSet<T, TAProperty> Include<TAProperty>(Expression<Func<T, TAProperty>> selector)
        {
            return new FetchedDbSet<T, TAProperty>(dbContext, fetchRequest.Fetch(selector));
        }

        public FetchedDbSet<T, TAProperty> IncludeMany<TAProperty>(Expression<Func<T, IEnumerable<TAProperty>>> selector)
        {
            return new FetchedDbSet<T, TAProperty>(dbContext, fetchRequest.FetchMany(selector));
        }

        public FetchedDbSet<T, TAProperty> ThenInclude<TAProperty>(Expression<Func<TProperty, TAProperty>> selector)
        {
            return new FetchedDbSet<T, TAProperty>(dbContext, fetchRequest.ThenFetch(selector));
        }

        public FetchedDbSet<T, TAProperty> ThenIncludeMany<TAProperty>(Expression<Func<TProperty,IEnumerable<TAProperty>>> selector)
        {
            return new FetchedDbSet<T, TAProperty>(dbContext, fetchRequest.ThenFetchMany(selector));
        }
    }
}